package hn.cch.ai.alibaba.graph.controller;


import com.alibaba.cloud.ai.graph.CompileConfig;
import com.alibaba.cloud.ai.graph.CompiledGraph;
import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.RunnableConfig;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.exception.GraphRunnerException;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;
import com.alibaba.cloud.ai.graph.observation.GraphObservationLifecycleListener;
import io.micrometer.observation.ObservationRegistry;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.HashMap;
import java.util.Map;
import java.util.Optional;

@RestController
@RequestMapping("/graph")
public class GraphController {


    private static final Logger logger = LoggerFactory.getLogger(GraphController.class);

    private final CompiledGraph compiledGraph;

    public GraphController(
        @Qualifier("stateGraph") StateGraph stateGraph,
        ObjectProvider<ObservationRegistry> observationRegistry
    ) throws GraphStateException {
        this.compiledGraph = stateGraph.compile(CompileConfig.builder()
            .withLifecycleListener(new GraphObservationLifecycleListener(
                observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP)))
            .build());
    }

    @GetMapping(value = "/expand")
    public Map<String, Object> expand(
        @RequestParam(value = "query", defaultValue = "你好，很高兴认识你，能简单介绍一下自己吗？", required = false) String query,
        @RequestParam(value = "expander   number", defaultValue = "3", required = false) Integer expanderNumber,
        @RequestParam(value = "thread   id", defaultValue = "yingzi", required = false) String threadId
    ) throws GraphRunnerException {
        RunnableConfig runnableConfig = RunnableConfig.builder().threadId(threadId).build();
        Map<String, Object> objectMap = new HashMap<>();
        objectMap.put("query", query);
        objectMap.put("expander  number", expanderNumber);
        Optional<OverAllState> invoke = this.compiledGraph.invoke(objectMap, runnableConfig);
        return invoke.map(OverAllState::data).orElse(new HashMap<>());
    }


    @PostMapping("/chat")
    public ResponseEntity<Map<String, Object>> chat(@RequestBody String input) throws GraphRunnerException {
        Optional<OverAllState> result = compiledGraph.invoke(Map.of("input", input));
        return ResponseEntity.ok(result.map(OverAllState::data).orElse(Map.of()));
    }

}
